import torch
from torch.utils.data import DataLoader, TensorDataset
import pickle
from torchvision import transforms
from PIL import Image

# https://blog.csdn.net/qq_43799400/article/details/122594950
# https://blog.csdn.net/qq_42902997/article/details/128667646

f = open('../data/mnist.pkl', 'rb')
training_data, validation_data, test_data = pickle.load(f, encoding='bytes')
f.close()

batch_size = 256

#  处理手写数字识别的数据集，返回三个迭代器
def load_data_wrapper():

    # 将训练数据转化成张量的形式
    trainX = torch.from_numpy(training_data[0])
    trainY = torch.from_numpy(training_data[1])

    # 将验证数据转化成张量的形式
    va_d_X = torch.from_numpy(validation_data[0])
    va_d_Y = torch.from_numpy(validation_data[1])

    # 将测试的数据转化成张量的形式
    test_dataX = torch.from_numpy(test_data[0])
    test_dataY = torch.from_numpy(test_data[1])

    # 通过TensorDatset函数将训练集，验证集和测试机的数据和标签进行组对
    train_set = TensorDataset(trainX, trainY)
    va_d_set = TensorDataset(va_d_X, va_d_Y)
    test_d_set = TensorDataset(test_dataX, test_dataY)

    # 通过DataLoader函数把三个数据集进行打包，并且设置包的规模，打乱样本顺序
    train_loader = DataLoader(dataset=train_set,batch_size=batch_size,shuffle=True)
    validation_loader = DataLoader(dataset=va_d_set,batch_size=batch_size,shuffle=False)
    test_loader = DataLoader(dataset=test_d_set,batch_size=batch_size,shuffle=False)

    return train_loader, validation_loader, test_loader

#  处理手写数字识别加法器的数据集，返回多个迭代器
def load_data_num_add():

    # 将训练数据转化成张量的形式
    trainX = torch.from_numpy(training_data[0])
    trainY = torch.from_numpy(training_data[1])

    # 将验证数据转化成张量的形式
    va_d_X = torch.from_numpy(validation_data[0])
    va_d_Y = torch.from_numpy(validation_data[1])

    # 将训练数据分成两部分
    trainX1 = trainX[:len(trainX) // 2]
    trainY1 = trainY[:len(trainY) // 2]
    trainX2 = trainX[len(trainX) // 2:]
    trainY2 = trainY[len(trainY) // 2:]

    # 将验证数据分成两部分
    va_d_X1 = va_d_X[:len(va_d_X) // 2]
    va_d_Y1 = va_d_Y[:len(va_d_Y) // 2]
    va_d_X2 = va_d_X[len(va_d_X) // 2:]
    va_d_Y2 = va_d_Y[len(va_d_Y) // 2:]

    # 创建数据集和数据加载器
    train_set1 = TensorDataset(trainX1, trainY1)
    train_loader1 = DataLoader(dataset=train_set1, batch_size=batch_size, shuffle=False)

    train_set2 = TensorDataset(trainX2, trainY2)
    train_loader2 = DataLoader(dataset=train_set2, batch_size=batch_size, shuffle=False)

    va_d_set1 = TensorDataset(va_d_X1, va_d_Y1)
    validation_loader1 = DataLoader(dataset=va_d_set1, batch_size=batch_size, shuffle=False)

    va_d_set2 = TensorDataset(va_d_X2, va_d_Y2)
    validation_loader2 = DataLoader(dataset=va_d_set2, batch_size=batch_size, shuffle=False)

    return train_loader1,train_loader2,validation_loader1,validation_loader2


# 加载手写数字图片，用于测试
def image_load():
    images_num = [7,2,1,0,4,1,4,9,5,9,0,6,9,0,1,5,9,7,3,4]
    img_dict={}
    for j in range(1, 21):
        path = './data/' + str(j) + '.jpg'
        image = Image.open(path)
        img= transforms.ToTensor()(image)
        img_and_num=[img,images_num[j-1]]
        img_dict[j]=img_and_num
    return img_dict
